import matplotlib.pyplot as plt
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler

# 加载数据  
data = pd.read_csv('Country-data.csv')  # 替换为你的数据文件路径

# 选择特征  
features = ['child_mort', 'health', 'income', 'life_expec', 'total_fer']
X = data[features]

# 数据标准化  
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# 应用 K-means 算法  
kmeans = KMeans(n_clusters=3)  # 假设我们选择 3 个集群作为示例  
kmeans.fit(X_scaled)

# 获取集群标签  
labels = kmeans.labels_

# 将标签添加到原始数据框中  
data['cluster'] = labels

# 查看每个集群的特征
print(data.groupby('cluster'))
# print(data.groupby('cluster').mean())

# 可视化（仅适用于二维或三维数据）  
# 例如，我们可以选择两个特征进行散点图可视化  
plt.scatter(X_scaled[:, 1], X_scaled[:, 0], c=labels, cmap='viridis')
plt.xlabel(features[1])
plt.ylabel(features[0])
plt.title('K-means Clustering')
plt.show()
